"""Various utility functions for terrain analysis.

- reset_nodata
- get_sun_position
- transform_coordinates
"""
import os
from typing import Optional, Union, Tuple

import numpy as np
from pysolar.solar import get_altitude, get_azimuth
import datetime
from osgeo import gdal, osr, ogr


def reset_nodata(
    geoimage: Union[str, gdal.Dataset],
    dst_nodata: Union[int, float],
    src_nodata: Optional[Union[int, float]] = None,
) -> gdal.Dataset:
    """Reset the nodata value of the input geoimage.

    Parameters
    ----------
    geoimage: str or gdal.Dataset
        Source geoimage GDAL dataset object or file path.
    dst_nodata: int or float
        Destination nodata value of the geoimage.
    src_nodata: int or float, optional
        Source nodata value of the geoimage.

    Returns
    -------
    ds: gdal.Dataset
        GDAL dataset of the geoimage after resetting the nodata value.
    """
    if not isinstance(geoimage, (str, gdal.Dataset)):
        raise TypeError(
            "Invalid type of 'geoimage'. Accepts a 'str' or a 'gdal.Dataset'."
        )
    if isinstance(geoimage, str) and not os.path.exists(geoimage):
        raise FileNotFoundError("Source geoimage file does not exists.")

    if isinstance(geoimage, gdal.Dataset):
        ds = geoimage
    else:
        ds = gdal.Open(geoimage, gdal.GA_Update)

    if ds is None:
        raise Exception("Invalid input gdal dataset or geoimage file.")

    for b in range(1, ds.RasterCount + 1):
        if src_nodata is None:
            src_nodata = ds.GetRasterBand(b).GetNoDataValue()
        band = ds.GetRasterBand(b).ReadAsArray()
        band = np.where(band == src_nodata, dst_nodata, band)
        ds.GetRasterBand(b).WriteArray(band)
        ds.GetRasterBand(b).SetNoDataValue(dst_nodata)

    if isinstance(geoimage, str):
        ds = None

    return ds


def get_sun_position(
    longitude: float, latitude: float, datetime_utc: datetime.datetime
) -> Tuple[float, float]:
    """Get the sun azimuth and altitude at the given location in space and time.

    Parameters
    ----------
    longitude: float
        The longitude of a location on earth between [-180, 180].
    latitude: float
        The latitude of a location on earth between [-90, 90].
    datetime_utc: datetime
        The date & time in UTC time zone.

    Returns
    -------
    azimuth: float
        Azimuth of the sun light, in degrees.
        0 if the light comes from the north, 90 from the east.
    altitude: float
        Altitude of the sun light, in degrees.
        90 if the light comes from above the DEM, 0 if it is from the horizon.
    """
    if not isinstance(datetime_utc, datetime.datetime):
        raise TypeError(
            "Invalid type of 'datetime_utc'. Accepts a datetime.datetime object."
        )
    azimuth = get_azimuth(latitude, longitude, datetime_utc)
    altitude = get_altitude(latitude, longitude, datetime_utc)

    return azimuth, altitude


def transform_coordinates(
    src_coords: Tuple[float, float], src_epsg_code: int, tgt_epsg_code: int
) -> Tuple[float, float]:
    """Transform a point coordinate (x, y) from the source coordinate system to the target one.

    Parameters
    ----------
    src_coords: Tuple(float, float)
        The coordinate of a point in the source coordinate system.
    src_epsg_code: int
        Source EPSG code.
    tgt_epsg_code: int
        Target EPSG code.

    Returns
    -------
    tgt_coords: Tuple(float, float)
        The coordinate of a point in the target coordinate system.
    """
    src_srs = osr.SpatialReference()
    src_srs.ImportFromEPSG(src_epsg_code)

    tgt_srs = osr.SpatialReference()
    tgt_srs.ImportFromEPSG(tgt_epsg_code)

    transform = osr.CoordinateTransformation(src_srs, tgt_srs)
    Point = ogr.Geometry(ogr.wkbPoint)
    Point.AddPoint(*src_coords)
    Point.Transform(transform)

    tgt_coords = (Point.GetX(), Point.GetY())
    Point = None

    return tgt_coords
